"""
ES和扩展F信号检测模型推理脚本
用于对新图片进行预测
"""

import torch
import torch.nn.functional as F
from PIL import Image
import numpy as np
import os
import argparse
from torchvision import transforms

from vit_model import create_model
from label_mapping import ID_TO_DESCRIPTION, DESCRIPTION_TO_ID

class ESFPredictor:
    def __init__(self, model_path, device=None):
        """
        初始化预测器
        
        Args:
            model_path: 训练好的模型路径
            device: 计算设备，默认自动选择
        """
        self.device = device if device else torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        
        # 加载模型
        print(f"加载模型: {model_path}")
        checkpoint = torch.load(model_path, map_location=self.device)
        
        # 创建模型
        self.model = create_model(num_classes=5)
        self.model.load_state_dict(checkpoint['model_state_dict'])
        self.model.to(self.device)
        self.model.eval()
        
        # 获取训练配置
        self.config = checkpoint.get('config', {})
        image_size = self.config.get('image_size', 224)
        
        # 图像预处理
        self.transform = transforms.Compose([
            transforms.Resize((image_size, image_size)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                               std=[0.229, 0.224, 0.225])
        ])
        
        print(f"模型加载完成，使用设备: {self.device}")
        print(f"模型验证准确率: {checkpoint.get('val_acc', 'N/A'):.2f}%")
    
    def preprocess_image(self, image_path):
        """
        预处理单张图片
        
        Args:
            image_path: 图片路径
            
        Returns:
            处理后的图片张量
        """
        try:
            # 加载图片
            image = Image.open(image_path).convert('RGB')
            
            # 应用变换
            image_tensor = self.transform(image).unsqueeze(0)  # 添加batch维度
            
            return image_tensor
        except Exception as e:
            print(f"处理图片 {image_path} 时出错: {e}")
            return None
    
    def predict_single(self, image_path, return_probs=False):
        """
        预测单张图片
        
        Args:
            image_path: 图片路径
            return_probs: 是否返回概率分布
            
        Returns:
            预测结果字典
        """
        # 预处理图片
        image_tensor = self.preprocess_image(image_path)
        if image_tensor is None:
            return None
        
        image_tensor = image_tensor.to(self.device)
        
        # 预测
        with torch.no_grad():
            outputs = self.model(image_tensor)
            probabilities = F.softmax(outputs, dim=1)
            confidence, predicted = torch.max(probabilities, 1)
            
            predicted_id = predicted.item()
            confidence_score = confidence.item()
            
            result = {
                'image_path': image_path,
                'predicted_class_id': predicted_id,
                'predicted_class': ID_TO_DESCRIPTION[predicted_id],
                'confidence': confidence_score
            }
            
            if return_probs:
                probs_dict = {}
                for i, desc in ID_TO_DESCRIPTION.items():
                    probs_dict[desc] = probabilities[0][i].item()
                result['probabilities'] = probs_dict
            
            return result
    
    def predict_batch(self, image_paths, return_probs=False):
        """
        批量预测多张图片
        
        Args:
            image_paths: 图片路径列表
            return_probs: 是否返回概率分布
            
        Returns:
            预测结果列表
        """
        results = []
        
        print(f"开始批量预测 {len(image_paths)} 张图片...")
        
        for i, image_path in enumerate(image_paths):
            if i % 100 == 0:
                print(f"处理进度: {i}/{len(image_paths)}")
            
            result = self.predict_single(image_path, return_probs)
            if result:
                results.append(result)
        
        print(f"批量预测完成，成功处理 {len(results)} 张图片")
        return results
    
    def predict_directory(self, directory_path, output_file=None, return_probs=False):
        """
        预测目录下所有图片
        
        Args:
            directory_path: 图片目录路径
            output_file: 输出文件路径（可选）
            return_probs: 是否返回概率分布
            
        Returns:
            预测结果列表
        """
        # 支持的图片格式
        image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif'}
        
        # 获取所有图片文件
        image_paths = []
        for filename in os.listdir(directory_path):
            if any(filename.lower().endswith(ext) for ext in image_extensions):
                image_paths.append(os.path.join(directory_path, filename))
        
        if not image_paths:
            print(f"在目录 {directory_path} 中未找到图片文件")
            return []
        
        print(f"找到 {len(image_paths)} 张图片")
        
        # 批量预测
        results = self.predict_batch(image_paths, return_probs)
        
        # 保存结果
        if output_file:
            self.save_results(results, output_file)
        
        return results
    
    def save_results(self, results, output_file):
        """
        保存预测结果到文件
        
        Args:
            results: 预测结果列表
            output_file: 输出文件路径
        """
        import json
        
        # 统计结果
        class_counts = {}
        for result in results:
            class_name = result['predicted_class']
            class_counts[class_name] = class_counts.get(class_name, 0) + 1
        
        # 准备输出数据
        output_data = {
            'total_images': len(results),
            'class_distribution': class_counts,
            'predictions': results
        }
        
        # 保存到JSON文件
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(output_data, f, indent=2, ensure_ascii=False)
        
        print(f"预测结果已保存到: {output_file}")
        
        # 打印统计信息
        print("\n预测结果统计:")
        for class_name, count in class_counts.items():
            percentage = count / len(results) * 100
            print(f"{class_name}: {count} 张 ({percentage:.1f}%)")

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='ES和扩展F信号检测模型推理')
    parser.add_argument('--model', required=True, help='模型文件路径')
    parser.add_argument('--image', help='单张图片路径')
    parser.add_argument('--directory', help='图片目录路径')
    parser.add_argument('--output', help='输出文件路径')
    parser.add_argument('--probs', action='store_true', help='输出概率分布')
    
    args = parser.parse_args()
    
    # 检查模型文件
    if not os.path.exists(args.model):
        print(f"模型文件不存在: {args.model}")
        return
    
    # 创建预测器
    predictor = ESFPredictor(args.model)
    
    if args.image:
        # 单张图片预测
        if not os.path.exists(args.image):
            print(f"图片文件不存在: {args.image}")
            return
        
        result = predictor.predict_single(args.image, return_probs=args.probs)
        if result:
            print(f"\n预测结果:")
            print(f"图片: {result['image_path']}")
            print(f"预测类别: {result['predicted_class']}")
            print(f"置信度: {result['confidence']:.4f}")
            
            if args.probs and 'probabilities' in result:
                print(f"\n各类别概率:")
                for class_name, prob in result['probabilities'].items():
                    print(f"  {class_name}: {prob:.4f}")
    
    elif args.directory:
        # 目录预测
        if not os.path.exists(args.directory):
            print(f"目录不存在: {args.directory}")
            return
        
        results = predictor.predict_directory(
            args.directory, 
            output_file=args.output,
            return_probs=args.probs
        )
        
        if results:
            print(f"\n预测完成，共处理 {len(results)} 张图片")
    
    else:
        print("请指定 --image 或 --directory 参数")

if __name__ == "__main__":
    main()